from pypinyin import lazy_pinyin, Style
from typing import List, Tuple
import re


def word_parser(sentence: str) -> Tuple[List[str], List[str]]:
    pinyins = lazy_pinyin(sentence)
    rhys = pinyin_parser(pinyins)
    rhys_tone = add_tone(sentence, rhys)
    return rhys, rhys_tone


def add_tone(sentence, ryhs):
    pinyins_tone = lazy_pinyin(sentence, style=Style.FINALS_TONE3)
    result = []
    for i, ryh in enumerate(ryhs):
        if pinyins_tone[i][-1] == '1' or pinyins_tone[i][-1] == '2' or pinyins_tone[i][-1] == '3' or pinyins_tone[i][
            -1] == '4':
            result.append(ryh + pinyins_tone[i][-1])
        else:
            result.append(ryh)
    return result


def pinyin_parser(pinyins: List[str]) -> List[List[str]]:
    parsed_pinyins = []
    for pinyin in pinyins:
        pinyin = pinyin_correction(pinyin)
        consonant, vowel = split_cv(pinyin)
        parsed_pinyins.append(weak_rhyme(vowel_parser(vowel)))
    return parsed_pinyins


def pinyin_correction(pinyin: str) -> str:
    # z/c/s + i -> z/c/s + I
    if re.match(r'[zcs]i$', pinyin):
        return pinyin.replace('i', 'I')

    # zh/ch/sh/r + i -> zh/ch/sh/r + II
    elif re.match(r'(?:zh|ch|sh|r)i$', pinyin):
        return pinyin.replace('i', 'II')

    # j/q/x/y + u/ue/un/uan -> j/q/x + v/ve/vn/van | v/ve/vn/van
    # yu == yi
    elif re.match(r'[jqxy]u', pinyin):
        if pinyin.endswith('u'):
            return re.sub(r'y*(.+)', r'\1', pinyin.replace('u', 'i'))
        return re.sub(r'y*(.+)', r'\1', pinyin.replace('u', 'v'))

    # y + a/e/ao/ou/an/ in/iang/ing/iong -> ia/ie/iao/iou/ian/ in/iang/ing/iong
    elif pinyin.startswith("y"):
        return re.sub(r'yi*(.*)', r'i\1', pinyin)

    # w + u/a/o/ai/ei/an/en/ang/eng        -> u/ua/uo/uai/uei/uan/uen/uang/ueng
    elif pinyin.startswith("w"):
        return re.sub(r'wu*(.*)', r'u\1', pinyin)

    # qiu -> qiou
    elif pinyin.endswith('iu'):
        return pinyin.replace('iu', 'iou')

    # cui -> cuei
    elif pinyin.endswith('ui'):
        return pinyin.replace('ui', 'uei')

    # lun -> luen
    elif pinyin.endswith('un'):
        return pinyin.replace('un', 'uen')

    return pinyin


def split_cv(pinyin):
    return re.findall(r'(ch|zh|sh|[^aeiIouv])*(.+)', pinyin)[0]


def vowel_parser(vowel: str) -> str:
    if vowel == 'van':
        return 'an'

    elif len(vowel) > 1 and vowel[0] == 'u':
        return vowel[1:]

    # for 'i', except 'in' and 'ing' (one vowel), 'ie' and 'ian' (sound different from 'e' and 'an')
    elif len(vowel) > 1 and vowel[0] == 'i' and vowel[1] != 'n' and vowel[1] != 'e' and vowel[1:] != 'an':
        return vowel[1:]

    return vowel


def weak_rhyme(pinyin):
    if pinyin.endswith('eng'):
        return pinyin.replace('eng', 'en')

    elif pinyin.endswith('ang'):
        return pinyin.replace('ang', 'an')

    elif pinyin.endswith('ing'):
        return pinyin.replace('ing', 'in')

    elif pinyin.endswith('ui'):
        return pinyin.replace('ui', 'ei')

    elif pinyin.endswith('ie'):
        return pinyin.replace('ui', 'ei')

    elif pinyin.endswith('ue'):
        return pinyin.replace('ue', 'ei')

    return pinyin
